package middleware

import (
	"github.com/gin-gonic/gin"
	"log"
	"net/http"
)

// 定义1GB的字节数常量，方便后续维护
const (
	KB = 1 << 10 // = 1024 Byte
	MB = 1 << 20
	GB = 1 << 30
)

// MaxBytesLimit 限制请求体总大小
func MaxBytesLimit(maxSize int64) gin.HandlerFunc {
	return func(ctx *gin.Context) {
		// 使用MaxBytesReader包装请求体，限制最大读取字节数
		ctx.Request.Body = http.MaxBytesReader(ctx.Writer, ctx.Request.Body, maxSize)

		contentType := ctx.GetHeader("Content-Type")
		if contentType == "multipart/form-data" {
			if err := ctx.Request.ParseMultipartForm(maxSize); err != nil {
				log.Printf("请求大小超过限制. %v", err.Error())
				ctx.String(http.StatusRequestEntityTooLarge, "请求大小超过限制 %d MB", maxSize/MB)
				ctx.Abort()
				return
			}
		}
		ctx.Next() // 继续处理请求
	}

}
